package com.shujia.flink.tf.process;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.HashMap;

public class Demo2KeyedProcessFunction {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStream<String> linesDS = env.socketTextStream("master", 8888);

        //按照班级分组
        KeyedStream<String, String> KeyByDS = linesDS.keyBy(line -> line.split(",")[4]);
        /*
         * process: 分组之后使用process可以再processElement中做聚合计算
         */
        KeyByDS.process(new KeyedProcessFunction<String, String, Tuple2<String, Integer>>() {

            //每一个task中共享同一个成员变量，不是每一个key体格num
            //private int num = 0;

            //再hashmap中为每一个班级都保存一个数量
            //每个task中共享同一个变量
            private HashMap<String, Integer> map = new HashMap<>();

            /**
             * processElement方法每一条数据执行一次
             * @param student：一行数据
             * @param ctx：上下文对象
             * @param out：用于将数据发送到下游
             */
            @Override
            public void processElement(String student,
                                       KeyedProcessFunction<String, String, Tuple2<String, Integer>>.Context ctx,
                                       Collector<Tuple2<String, Integer>> out) throws Exception {
                System.out.println(map);

                //获取当前分组的key
                String clazz = ctx.getCurrentKey();

                //第一次没有返回0, 后面就返回对应的数量
                Integer num = map.getOrDefault(clazz, 0);

                //累加计算
                num++;

                //将累加结果保存到map中
                map.put(clazz, num);

                //将结果发送到下游
                out.collect(Tuple2.of(clazz, num));
            }
        }).print();

        env.execute();

    }
}
